import cv2
import numpy as np
import signal
import sys
import pytesseract
import time
import psutil
import threading
import queue

# 性能优化参数
REDUCE_RESOLUTION = True        # 降低识别区域分辨率
SKIP_FRAMES = 0                 # 每X帧处理1次，0表示处理所有帧
MIN_CONFIDENCE = 30             # 最小识别置信度(值越高越严格)
CHECK_MEMORY = False            # 禁用内存监控
USE_MULTITHREADING = True       # 启用多线程处理

# 屏幕分辨率设置
DISPLAY_WIDTH = 480
DISPLAY_HEIGHT = 800
PROCESS_WIDTH = 320             # 处理分辨率
PROCESS_HEIGHT = 240            # 处理分辨率

# 按钮设置
BUTTON_WIDTH = 200
BUTTON_HEIGHT = 80
BUTTON_COLOR = (0, 0, 255)
BUTTON_COLOR_PRESSED = (0, 100, 255)
BUTTON_TEXT = "EXIT"
BUTTON_TEXT_COLOR = (255, 255, 255)

# 识别区域设置（正方形，中央）
ROI_SIZE_RATIO = 0.5            # 识别区域占屏幕宽度的比例
ROI_BORDER_THICKNESS = 2        # 识别区域边框厚度

# OCR优化参数
ADAPTIVE_BLOCK_SIZE = 11        # 自适应阈值块大小(必须为奇数)
ADAPTIVE_C = 2                  # 自适应阈值常量
MORPH_KERNEL_SIZE = 3           # 形态学操作内核大小
DILATE_ITERATIONS = 1           # 膨胀迭代次数
ERODE_ITERATIONS = 1            # 腐蚀迭代次数
CONTOUR_AREA_MIN = 100          # 过滤小轮廓的最小面积
CONTOUR_AREA_MAX = 5000         # 过滤大轮廓的最大面积

# 内存监控变量（已禁用）
process = psutil.Process()
max_memory = 0
max_memory_percent = 0.0

# 创建线程安全队列
frame_queue = queue.Queue(maxsize=3)  # 限制队列大小防止内存溢出
result_queue = queue.Queue(maxsize=1)

# 图像处理和OCR识别线程
def processing_thread():
    global last_recognized
    
    while True:
        try:
            # 获取待处理帧
            frame, roi_rect = frame_queue.get(timeout=1)
            
            # 退出信号
            if frame is None:
                break
                
            # 提取识别区域（正方形）
            roi_x, roi_y, roi_size = roi_rect
            roi = frame[roi_y:roi_y+roi_size, roi_x:roi_x+roi_size]
            
            # 预处理图像以提高OCR准确率
            processed = preprocess_image(roi)
            
            # 使用增强配置识别数字
            recognized_digits = recognize_digits(processed)
            
            # 将结果放入队列(覆盖旧结果)
            try:
                result_queue.get_nowait()  # 清空队列
            except queue.Empty:
                pass
            result_queue.put(recognized_digits)
            
            frame_queue.task_done()
        except queue.Empty:
            continue
        except Exception as e:
            print(f"处理线程错误: {e}")

def signal_handler(sig, frame):
    """信号处理函数，捕获Ctrl+C"""
    print("\n程序停止")
    
    # 清理线程
    if USE_MULTITHREADING:
        # 清空队列并发送退出信号
        try:
            while True:
                frame_queue.get_nowait()
        except queue.Empty:
            pass
        frame_queue.put((None, None))  # 退出信号
    
    sys.exit(0)

def draw_exit_button(frame, pressed=False):
    """在画面上绘制退出按钮"""
    x = frame.shape[1] - BUTTON_WIDTH - 20
    y = frame.shape[0] - BUTTON_HEIGHT - 20
    color = BUTTON_COLOR_PRESSED if pressed else BUTTON_COLOR
    cv2.rectangle(frame, (x, y), (x + BUTTON_WIDTH, y + BUTTON_HEIGHT), color, -1)
    text_size = cv2.getTextSize(BUTTON_TEXT, cv2.FONT_HERSHEY_SIMPLEX, 1.2, 2)[0]
    text_x = x + (BUTTON_WIDTH - text_size[0]) // 2
    text_y = y + (BUTTON_HEIGHT + text_size[1]) // 2
    cv2.putText(frame, BUTTON_TEXT, (text_x, text_y), 
                cv2.FONT_HERSHEY_SIMPLEX, 1.2, BUTTON_TEXT_COLOR, 2)
    return frame, (x, y, BUTTON_WIDTH, BUTTON_HEIGHT)

def preprocess_image(img):
    """增强图像预处理，提高OCR识别准确率"""
    # 转换为灰度图
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    
    # 高斯模糊减少噪声
    blur = cv2.GaussianBlur(gray, (5, 5), 0)
    
    # 自适应阈值处理，优化参数
    thresh = cv2.adaptiveThreshold(
        blur, 255, 
        cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
        cv2.THRESH_BINARY_INV, 
        ADAPTIVE_BLOCK_SIZE, 
        ADAPTIVE_C
    )
    
    # 形态学操作清理小噪声
    kernel = np.ones((MORPH_KERNEL_SIZE, MORPH_KERNEL_SIZE), np.uint8)
    opening = cv2.morphologyEx(thresh, cv2.MORPH_OPEN, kernel, iterations=1)
    
    # 膨胀增强数字
    dilated = cv2.dilate(opening, kernel, iterations=DILATE_ITERATIONS)
    
    # 腐蚀调整数字粗细
    eroded = cv2.erode(dilated, kernel, iterations=ERODE_ITERATIONS)
    
    # 轮廓检测和过滤，去除小噪点
    contours, _ = cv2.findContours(eroded.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    mask = np.zeros_like(eroded)
    
    for contour in contours:
        area = cv2.contourArea(contour)
        if CONTOUR_AREA_MIN < area < CONTOUR_AREA_MAX:
            cv2.drawContours(mask, [contour], -1, 255, -1)
    
    # 应用掩码到原始阈值图像
    result = cv2.bitwise_and(eroded, eroded, mask=mask)
    
    return result

def recognize_digits(img):
    """增强数字识别，带置信度过滤"""
    global last_recognized
    
    try:
        # 调整图像大小到Tesseract最佳识别尺寸(高度约30-50像素)
        h, w = img.shape[:2]
        if h < 30:
            scale = 30 / h
            img = cv2.resize(img, None, fx=scale, fy=scale)
        elif h > 50:
            scale = 50 / h
            img = cv2.resize(img, None, fx=scale, fy=scale)
        
        # 反转图像(Tesseract对浅色背景上的深色文本处理更好)
        inverted = cv2.bitwise_not(img)
        
        # 使用Tesseract LSTM引擎，优化参数
        config = (
            f'--psm 10 --oem 3 '  # PSM 10 = 将图像视为单个字符
            f'-c tessedit_char_whitelist=0123456789 '
            f'-c min_char_conf={MIN_CONFIDENCE} '
            f'-c load_system_dawg=false '
            f'-c load_freq_dawg=false'
        )
        
        # 获取带置信度的OCR结果
        result = pytesseract.image_to_string(inverted, config=config).strip()
        
        # 验证结果
        if result and result.isdigit():
            last_recognized = result
            return result
        
        return last_recognized if 'last_recognized' in globals() else ""
    
    except Exception as e:
        print(f"识别错误: {e}")
        return last_recognized if 'last_recognized' in globals() else ""

def mouse_callback(event, x, y, flags, param):
    """鼠标/触摸事件回调函数"""
    global button_pressed, exit_requested
    button_rect = param
    
    if event == cv2.EVENT_LBUTTONDOWN:
        bx, by, bw, bh = button_rect
        if bx <= x <= bx + bw and by <= y <= by + bh:
            button_pressed = True
    
    elif event == cv2.EVENT_LBUTTONUP:
        button_pressed = False
        bx, by, bw, bh = button_rect
        if bx <= x <= bx + bw and by <= y <= by + bh:
            exit_requested = True

def main():
    global exit_requested, show_roi, show_result, frame_counter, button_pressed, last_recognized, USE_MULTITHREADING
    
    # 初始化所有全局变量
    frame_counter = 0
    button_pressed = False
    exit_requested = False
    show_roi = True
    show_result = True
    last_recognized = ""
    
    signal.signal(signal.SIGINT, signal_handler)
    
    # 创建窗口
    cv2.namedWindow("Camera Feed", cv2.WINDOW_NORMAL)
    cv2.resizeWindow("Camera Feed", DISPLAY_WIDTH, DISPLAY_HEIGHT)
    
    # 打开默认摄像头
    cap = cv2.VideoCapture(0)
    
    # 设置摄像头分辨率
    cap.set(cv2.CAP_PROP_FRAME_WIDTH, DISPLAY_WIDTH)
    cap.set(cv2.CAP_PROP_FRAME_HEIGHT, DISPLAY_HEIGHT)
    
    # 尝试设置摄像头对焦模式(如果支持)
    try:
        cap.set(cv2.CAP_PROP_AUTOFOCUS, 0)  # 禁用自动对焦
        cap.set(cv2.CAP_PROP_FOCUS, 0)      # 设置最小对焦距离
    except:
        print("无法设置摄像头对焦参数")
    
    if not cap.isOpened():
        print("无法打开摄像头")
        sys.exit(1)
    
    print("CPU优化版数字识别程序已启动")
    print("触摸右下角按钮或按 'q' 键退出")
    print("按 'r' 键切换识别区域显示")
    print("按 's' 键切换识别结果显示")
    print("按 'm' 键切换多线程模式")
    
    # 初始化按钮区域
    button_rect = (0, 0, BUTTON_WIDTH, BUTTON_HEIGHT)
    
    # 设置鼠标回调函数
    cv2.setMouseCallback("Camera Feed", mouse_callback, button_rect)
    
    # 启动处理线程
    if USE_MULTITHREADING:
        process_thread = threading.Thread(target=processing_thread, daemon=True)
        process_thread.start()
    
    try:
        while True:
            start_time = time.time()
            
            # 读取一帧
            ret, frame = cap.read()
            if not ret:
                print("无法接收帧，退出...")
                break
            
            frame_counter += 1
            
            # 降低处理分辨率
            if REDUCE_RESOLUTION:
                process_frame = cv2.resize(frame, (PROCESS_WIDTH, PROCESS_HEIGHT))
            else:
                process_frame = frame.copy()
            
            # 准备显示帧
            display_frame = frame.copy()
            
            # 计算识别区域大小和位置(正方形，居中)
            if REDUCE_RESOLUTION:
                # 处理分辨率下的正方形识别区域
                roi_size = int(min(PROCESS_WIDTH, PROCESS_HEIGHT) * ROI_SIZE_RATIO)
                roi_x = (PROCESS_WIDTH - roi_size) // 2
                roi_y = (PROCESS_HEIGHT - roi_size) // 2
                
                # 转换到显示分辨率的坐标
                scale_x = DISPLAY_WIDTH / PROCESS_WIDTH
                scale_y = DISPLAY_HEIGHT / PROCESS_HEIGHT
                display_roi_x = int(roi_x * scale_x)
                display_roi_y = int(roi_y * scale_y)
                display_roi_size = int(roi_size * scale_x)  # 保持正方形比例
            else:
                # 显示分辨率下的正方形识别区域
                roi_size = int(min(DISPLAY_WIDTH, DISPLAY_HEIGHT) * ROI_SIZE_RATIO)
                roi_x = (DISPLAY_WIDTH - roi_size) // 2
                roi_y = (DISPLAY_HEIGHT - roi_size) // 2
                display_roi_x, display_roi_y, display_roi_size = roi_x, roi_y, roi_size
            
            # 显示识别区域
            if show_roi:
                cv2.rectangle(display_frame, (display_roi_x, display_roi_y), 
                            (display_roi_x+display_roi_size, display_roi_y+display_roi_size), 
                            (0, 255, 0), ROI_BORDER_THICKNESS)
            
            # 性能优化:跳帧处理
            if frame_counter % (SKIP_FRAMES + 1) == 0:
                # 多线程模式:将帧放入队列供处理线程使用
                if USE_MULTITHREADING:
                    try:
                        # 如果队列已满，丢弃旧帧以保持最新帧
                        if frame_queue.full():
                            frame_queue.get_nowait()
                        frame_queue.put_nowait((process_frame.copy(), (roi_x, roi_y, roi_size)))
                    except Exception as e:
                        print(f"队列错误: {e}")
                
                # 单线程模式:直接在主线程处理
                else:
                    # 提取识别区域
                    roi = process_frame[roi_y:roi_y+roi_size, roi_x:roi_x+roi_size]
                    
                    # 预处理图像
                    processed = preprocess_image(roi)
                    
                    # 识别数字
                    recognized_digits = recognize_digits(processed)
                    
                    # 显示识别结果
                    if show_result and recognized_digits:
                        cv2.putText(display_frame, f"Recognized: {recognized_digits}", 
                                  (30, 60), cv2.FONT_HERSHEY_SIMPLEX, 1.5, 
                                  (0, 255, 0), 2)
            
            # 从结果队列获取最新识别结果
            if USE_MULTITHREADING and show_result:
                try:
                    recognized_digits = result_queue.get_nowait()
                    if recognized_digits:
                        cv2.putText(display_frame, f"Recognized: {recognized_digits}", 
                                  (30, 60), cv2.FONT_HERSHEY_SIMPLEX, 1.5, 
                                  (0, 255, 0), 2)
                except queue.Empty:
                    pass
            
            # 显示CPU占用率
            cpu_text = f"CPU: {psutil.cpu_percent(interval=None):.1f}%"
            cv2.putText(display_frame, cpu_text, (30, 110), 
                       cv2.FONT_HERSHEY_SIMPLEX, 1.2, 
                       (0, 255, 0), 2)
            
            # 绘制退出按钮
            display_frame, button_rect = draw_exit_button(display_frame, button_pressed)
            cv2.setMouseCallback("Camera Feed", mouse_callback, button_rect)
            
            # 显示帧
            cv2.imshow("Camera Feed", display_frame)
            
            # 检查退出
            key = cv2.waitKey(1) & 0xFF
            if key == ord('q') or exit_requested:
                break
            elif key == ord('r'):
                show_roi = not show_roi
                print(f"Recognition area display: {'ON' if show_roi else 'OFF'}")
            elif key == ord('s'):
                show_result = not show_result
                print(f"Recognition result display: {'ON' if show_result else 'OFF'}")
            elif key == ord('m'):  # 切换多线程模式
                if USE_MULTITHREADING:
                    # 禁用多线程
                    USE_MULTITHREADING = False
                    print("Multithreading mode disabled")
                else:
                    # 启用多线程
                    USE_MULTITHREADING = True
                    process_thread = threading.Thread(target=processing_thread, daemon=True)
                    process_thread.start()
                    print("Multithreading mode enabled")
            
            if cv2.getWindowProperty("Camera Feed", cv2.WND_PROP_VISIBLE) < 1:
                break
                
    except KeyboardInterrupt:
        print("\n程序被用户中断")
        
    finally:
        # 释放资源
        cap.release()
        cv2.destroyAllWindows()
        
        # 等待处理线程结束
        if USE_MULTITHREADING:
            try:
                # 清空队列并发送退出信号
                while True:
                    frame_queue.get_nowait()
            except queue.Empty:
                pass
            frame_queue.put((None, None))  # 退出信号
        
        print("资源已释放")

if __name__ == "__main__":
    main()
    